import numpy as np
import cv2 as cv

class MixUpWithMaskConvert(object):
    
    def __init__(self, is_random = True):
        self.is_random = is_random
        self.__r = 0.2
    
    def setR(self, r):
        self.__r = r    
    
    # 转换
    def convert(self, imgs, masks):
        r = self.__r
        if self.is_random:
            r = np.random.beta(0.2, 0.2)
            
        img = imgs[0]
        mask = masks[0]
        for i in range(1, len(imgs)):
            xi = img
            xj = imgs[i]
            # 融合背景
            img = r * xi + (1 - r) * xj
            
            xi_mask = mask
            xj_mask = masks[i]
            
            xi_mask_content = xi_mask > 0
            xj_mask_content = xj_mask > 0
            
            img[xi_mask_content] = xi[xi_mask_content]
            img[xj_mask_content] = xj[xj_mask_content]
            
            mask = np.maximum(xi_mask, xj_mask)
            
        return img, mask
        


    